import torch
import torch.nn as nn
import numpy as np

np.random.seed(42)
torch.manual_seed(42)

class UltimatePINN(nn.Module):
    def __init__(self, c_mean, c_std):
        super().__init__()
        self.c_mean = c_mean
        self.c_std = c_std
        self.net = nn.Sequential(
            nn.Linear(1, 32), 
            nn.Tanh(), 
            nn.Linear(32, 32), 
            nn.Tanh(), 
            nn.Linear(32, 1)
        )
        self.log_Ct = nn.Parameter(torch.log(torch.tensor([48.0])))
        self.log_AUC = nn.Parameter(torch.log(torch.tensor([34.0])))
        self.log_lb = nn.Parameter(torch.log(torch.tensor([2.4])))
        self.Cd = nn.Parameter(torch.tensor([-15.0])) 

    @property
    def params(self):
        return torch.exp(self.log_Ct), torch.exp(self.log_AUC), torch.exp(self.log_lb), self.Cd

    def forward(self, t):
        Ct_val = torch.exp(self.log_Ct)
        Ct_scaled = (Ct_val - self.c_mean) / self.c_std
        return Ct_scaled + (1.0 - torch.exp(-t)) * self.net(t)

    def physics_loss(self, t):
        C = self.forward(t) * self.c_std + self.c_mean
        Ct, AUC, lb, Cd = self.params
        dCdt = torch.autograd.grad(C, t, torch.ones_like(C), create_graph=True)[0]
        rhs = 0.5 * AUC * lb**3 * (2*t - lb*t**2) * torch.exp(-lb*t) + Cd * lb * torch.exp(-lb*t)
        return torch.mean((dCdt - rhs) ** 2)

def train_core(t_raw, c_raw):
    c_mean, c_std = c_raw.mean(), c_raw.std()
    
    t_data = torch.tensor(t_raw, dtype=torch.float32).view(-1, 1).requires_grad_(True)
    c_data = torch.tensor((c_raw - c_mean) / c_std, dtype=torch.float32).view(-1, 1)
    t_phys = torch.linspace(0, 7, 100).view(-1, 1).requires_grad_(True)
    
    model = UltimatePINN(c_mean, c_std)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    for _ in range(8000):
        optimizer.zero_grad()
        loss_data = torch.mean((model(t_data) - c_data) ** 2)
        loss_phys = model.physics_loss(t_phys) / c_std ** 2
        (loss_data + 1.5 * loss_phys).backward()
        optimizer.step()
        
    lbfgs = torch.optim.LBFGS(model.parameters(), max_iter=200, line_search_fn="strong_wolfe")
    
    def closure():
        lbfgs.zero_grad()
        loss = torch.mean((model(t_data) - c_data)**2) + 1.0 * model.physics_loss(t_phys) / c_std**2
        loss.backward()
        return loss
        
    lbfgs.step(closure)
    return model